# -*- coding: utf-8 -*-
"""
Created on xxx  2024-2025
PyTorch编程实现循环神经网络的上证指数预测
数据集：stockPrice.csv（2024-5到2024-11月份上证指数）
实验环境：python3.9.13、torch1.10.2+cu113
数据处理：需要做归一化
神经网络：lstm + linear
实验目的：循环神经网络流程及网络结构（重要的是数据维度和LSTM参数维度区别！）
实验要求：拟合图、y轴数值还原、loss损失画图
@author: htq
"""
import pandas as pd
import numpy as np
import torch.nn as nn
import torch
import matplotlib.pyplot as plt

import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'  # 允许重复加载动态链接库

plt.rcParams['figure.figsize'] = (8.0, 6.0)  # 设置画图窗口大小
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体，以正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 设置正常显示符号

# ==============数据处理==============
data_csv = pd.read_csv('./stockPrice.csv', usecols=[4], encoding='gbk')  # 读取csv文件第4列收盘价
dataset = data_csv.values  # 只获取值，不要索引
dataset = dataset.astype('float32')  # 统一为pytorch常用精度
close_price = dataset  # 拿到没有归一化的真实值
# (x - min)/(max - min) 最大最小值归一化 --> [0-1]
max_value = np.max(dataset)
min_value = np.min(dataset)
scale = max_value - min_value
dataset = list(map(lambda x: (x - min_value) / scale, dataset))

# 1 2 3 4 5
# 1 2 -> 3
# 2 3 -> 4
# 3 4 -> 5

# 构建数据集  昨天 今天 --> 明天  [3302, 3303] --> [3304]
look_back = 2  # 2天进行预测


def create_dataset(dataset):
    data_X, data_Y = [], []
    for i in range(len(dataset) - look_back):
        data_X.append(dataset[i:i + look_back])
        data_Y.append(dataset[i + look_back])
    return np.array(data_X), np.array(data_Y)


data_X, data_Y = create_dataset(dataset)
train_size = int(len(data_X) * 1)  # 划分数据集，训练集，测试集， 82法则， 0.8 0.2
train_X = data_X[:train_size]
train_Y = data_Y[:train_size]
test_X = data_X[-1]

train_x = train_X.reshape(-1, 1, look_back)  # 变成三维（seq_len, batch_size, input_size）
train_y = train_Y.reshape(-1, 1, 1)
test_x = test_X.reshape(-1, 1, look_back)

train_x = torch.from_numpy(train_x)  # 变成tensor类型
train_y = torch.from_numpy(train_y)
test_x = torch.from_numpy(test_x)


# ==============构建网络==============
class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=2, output_size=1):
        super(LSTM, self).__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x, _ = self.lstm(x)
        s, b, h = x.shape
        x = x.view(s * b, h)  # 改变成二维形状
        x = self.linear(x)
        x = x.view(s, b, -1)  # 改变成三维形式方便后面绘图使用，也可不改变！
        return x


# ==============训练网络==============
net = LSTM(look_back, 4)
cirterion = nn.MSELoss()  # 定义损失函数，回归问题
optimizer = torch.optim.Adam(net.parameters(), lr=0.01)  # 定义优化器，传入模型参数和学习率

loss_list = []
for i in range(1000):
    out = net(train_x)  # 前向传播
    loss = cirterion(out, train_y)  # 计算损失
    optimizer.zero_grad()  # 梯度清零
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数
    loss_list.append(loss.detach().numpy())  # 保存损失
    if i % 100 == 0:
        print('Epoch:{}, Loss:{:.4f}'.format(i, loss))  # 每100轮打印损失值

# ==============性能评估==============
net = net.eval()  # 开启评估模式
pred_test = net(train_x)
pred_test = pred_test.view(-1).data.numpy()  # 改变3维为1维，为了plt画图

for index, value in enumerate(pred_test):  # 数据从0-1还原
    pred_test[index] = value * scale + min_value

plt.subplot(2, 1, 1)  # 创建一个画图面板，2行，1列,当前是第一个子图
plt.suptitle("2024年上证指数预测", fontsize=14, fontweight='bold')
plt.plot(pred_test, 'r', label='预测值')
plt.plot(close_price[2:], 'b', label='真实值')
plt.xlabel('2024年5月-2024年11月', fontsize=14, fontweight='bold')
plt.ylabel('上证指数预测值', fontsize=14, fontweight='bold')
plt.legend(loc='best')  # 自动确定图例在图形中的最佳位置

plt.subplots_adjust(hspace=0.4)  # 0.4表示子图之间垂直距离为子图宽度的40%

plt.subplot(2, 1, 2)  # 创建一个画图面板，2行，1列,当前是第二个子图
plt.plot(loss_list, 'g', label='损失值')
plt.xlabel('迭代数', fontsize=14, fontweight='bold')
plt.ylabel('损失值', fontsize=14, fontweight='bold')
plt.legend(loc='upper right')  # 确定图例在图形中的右上角
plt.show()

# ==============预测明天==============
pred_tom = net(test_x)
pred_tom = pred_tom.view(-1).data.numpy()  # 改变3维为1维，为了plt画图
for index, value in enumerate(pred_tom):  # 数据从0-1还原
    pred_tom[index] = value * scale + min_value
print(round(pred_tom.item(), 2))

"""
LSTM 的输入维度 (seq_len, batch_size, input_size) 和定义 LSTM 网络时的参数 input_size、hidden_size 和 num_layers 
是不同的概念，它们分别代表不同的内容（前者 描述了输入数据的形状，后者描述了 LSTM 网络的结构）

LSTM 输入维度 (seq_len, batch_size, input_size)
seq_len: 序列的长度，即每个样本中的时间步数或序列长度。
batch_size: 批处理大小，即一次训练中使用的样本数量。
input_size: 每个时间步的特征数量，即每个输入向量的维度。
这个维度是用于描述输入数据的形状，它决定了每个时间步的输入特征数量。

LSTM 网络参数 (input_size, hidden_size, num_layers)
input_size: 这是每个时间步输入特征的维度，即每个输入向量的大小。例如，如果每个时间步有5个特征，那么 input_size 就是5。
hidden_size: 这是隐藏层的特征维度，即 LSTM 单元在每个时间步输出的向量大小。这个值决定了 LSTM 单元内部状态的维度。
num_layers: 这是 LSTM 的层数。多层 LSTM 可以堆叠在一起，每一层都会接收前一层的输出作为输入。通过增加层数，模型可以捕捉更复杂的时序依赖关系。
"""
